package com.sunzm.spark.structstream.state

import org.apache.spark.sql.streaming._
import org.apache.spark.sql.{DataFrame, Dataset, SparkSession}
import org.slf4j.{Logger, LoggerFactory}

/**
 *
 * 使用flatMapGroupsWithState模拟单词计数
 *
 * @author Administrator
 * @version 1.0
 * @date 2021-07-05 0:38
 */
object StreamingFlatMapGroupsWithStateDemo {
  private val logger: Logger = LoggerFactory.getLogger(this.getClass.getName.stripSuffix("$"))

  def main(args: Array[String]): Unit = {
    val host = "82.156.210.70"
    val port = 9999

    val spark: SparkSession = SparkSession
      .builder()
      .appName(this.getClass.getSimpleName.stripSuffix("$"))
      .master("local[*]")
      .config("spark.default.parallelism", 6)
      .config("spark.sql.shuffle.partitions", 6)
      .getOrCreate()

    import spark.implicits._

    //从一个socket端口创建一个DataFrame
    val lines: DataFrame = spark.readStream
      .format("socket")
      .option("host", host)
      .option("port", port)
      .load()

    // Split the lines into words, treat words as sessionId of events
    //hello,spark,hello
    val events: Dataset[Event] = lines
      .as[String]
      .flatMap(_.split(",")
        .map(word => Event(word, 1)))

    /**
     * /**
     * * 第一个泛型：Key的类型
     * * 第二个泛型：Value的类型
     * * 第三个泛型：State（状态的类型）
     * * 第四个类型：返回值的类型 (输出的数据类型)
     * */
     */
    val resultDS: Dataset[String] = events
      .groupByKey(event => event.word)
      .flatMapGroupsWithState(OutputMode.Update(), GroupStateTimeout.NoTimeout())(flatMapFunc)

    // Start running the query that prints the session updates to the console
    val query = resultDS
      .writeStream
      .outputMode(OutputMode.Update())
      .format("console")
      .start()

    query.awaitTermination()

  }

  def flatMapFunc(key: String, values: Iterator[Event], state: GroupState[WordCountState]): Iterator[String] = {

    val valueSize: Int = values.size

    val newWordCount: Long = if (state.exists) {
      val oldWordCount: WordCountState = state.get
      oldWordCount.count + valueSize
    } else {
      //第一次
      valueSize
    }

    val newWordCountState: WordCountState = WordCountState(key, newWordCount, System.currentTimeMillis())

    //更新状态
    state.update(newWordCountState)

    //清空状态
    //state.remove()
    //设置状态多久后超时
    //state.setTimeoutDuration("10 seconds")

    Iterator[String](s"${newWordCountState.word}, 当前累计出现次数: ${newWordCountState.count}")
  }

  /** User-defined data type representing the input events */
  private case class Event(word: String, count: Long)

  private case class WordCountState(word: String, count: Long, timestamp: Long)
}

